package scu.maqiang.numeric;

import java.util.Arrays;

public class TridiagMatrix {
	public TridiagMatrix(int n) {
		this.n = n;
		diag = new double[n];
		lower = new double[n - 1];
		upper = lower.clone();
	}
	
	public TridiagMatrix(SRMatrix A) {
		this(A.getM());
		int lowIdx = 0;
		int upperIdx = 0;
		for(int i = 0; i < n; i++) {
			for(SREntry ele : A.mat[i]) {
				int diff = ele.col - i; 
				if (diff == -1) {
					lower[lowIdx++] = ele.val;
				}else if (diff == 0) {
					diag[i] = ele.val;
				}else if (diff == 1) {
					upper[upperIdx++] = ele.val;
				}
			}
		}
	}
	
	public TridiagMatrix(double[][] A) {
		this(A.length);
		diag[0] = A[0][0];
		upper[0] = A[0][1];
		lower[n - 2] = A[n - 1][n - 2];
		diag[n - 1] = A[n - 1][n - 1];
		for(int i = 1; i < n - 1; i++) {
			lower[i - 1] = A[i][i - 1];
			diag[i] = A[i][i];
			upper[i] = A[i][i + 1];
		}
	}
	
	public void LUFactorize() {
		int n = diag.length;
		for(int i = 0; i < n - 1; i++) {
			lower[i] /= diag[i];
			diag[i + 1] -= upper[i] * lower[i];
		}
	}
	
	public void LUSolve(double[] b, double[] x) {
		int n = diag.length;
		for(int i = 1; i < n; i++) {
			b[i] -= lower[i - 1] * b[ i - 1];
		}
		b[n - 1] /= diag[n - 1];
		for(int i = n - 2; i >= 0; i--) {
			b[i] = (b[i] - b[i + 1] * upper[i])  / diag[i];
		}
		Arrays.setAll(x, i -> b[i]);
	}
	
	public static void main(String[] args) {
		double[][] A = {{3, 1, 0, 0}, {1, 3, 1, 0}, {0, 1, 3, 1}, {0, 0, 1, 2}};
		TridiagMatrix triMat = new TridiagMatrix(A);
		System.out.println(Arrays.toString(triMat.diag));
		System.out.println(Arrays.toString(triMat.lower));
		System.out.println(Arrays.toString(triMat.upper));
		double[] b = {4, 5, 5, 4};
		double[] x = new double[4];
		triMat.LUFactorize();
		triMat.LUSolve(b, x);
		System.out.println(Arrays.toString(x));
	}
	
	int n;
	double[] diag;
	double[] lower;
	double[] upper;
}
